'''
Author: Morphlng
Date: 2022-05-05 18:19:48
LastEditTime: 2022-05-06 16:41:35
LastEditors: Morphlng
Description: 
FilePath: \flappy_bird\flappyBird_model.py
'''

import paddle.nn as nn
import paddle.nn.functional as F
import parl


class FlappyBirdModel(parl.Model):
    """ Linear network to solve FlappyBird-V0 problem.

    Args:
        obs_dim (int): Dimension of observation space.
        act_dim (int): Dimension of action space.
    """

    def __init__(self, obs_dim, act_dim):
        super(FlappyBirdModel, self).__init__()
        hid1_size = 64
        hid2_size = 32
        # 3 fully connected layers
        self.fc1 = nn.Linear(obs_dim, hid1_size)
        self.fc2 = nn.Linear(hid1_size, hid2_size)
        self.fc3 = nn.Linear(hid2_size, act_dim)

    def forward(self, x):
        # input the state, output corresponding action
        h1 = F.relu(self.fc1(x))
        h2 = F.relu(self.fc2(h1))
        Q = self.fc3(h2)
        return Q
